import pandas as pd
from typing import Optional


class Context:
    """
    一个安全的数据访问代理，是策略与历史数据交互的唯一接口。

    它的核心职责是扮演一个“时间旅行安全气囊”，确保在回测的任何给定
    时间点，策略代码都只能访问到该时间点之前（不包括该时间点）的
    市场数据，从而从根本上杜絕未来函数。

    它被设计为轻量级对象，可以在回测引擎的主循环中高效地被反复创建。
    """

    def __init__(self, full_data: pd.DataFrame, current_time: pd.Timestamp):
        """
        初始化Context。

        Args:
            full_data (pd.DataFrame): 包含了所有K线和预计算指标的完整主DataFrame。
            current_time (pd.Timestamp): 当前回测的时间点。
                                       Context将只提供严格早于此时间的数据。
        """
        self.current_time = current_time

        # 【核心安全机制】
        # 在对象创建时，就立即对完整数据进行一次性的历史切片。
        # 后续所有的数据访问都将基于这个被“净化”过的历史数据视图。
        # 我们使用严格小于 (<) 来确保当前时间点的数据对策略是不可见的，
        # 模拟了策略在当前K线形成过程中进行决策的真实场景。
        self._historical_data = full_data.loc[full_data.index < current_time]

    def get_latest_bar(self) -> Optional[pd.Series]:
        """
        获取最新的、对策略完全可见的K线数据（及其所有指标）。

        这是策略组件中最常用的数据访问方法。

        Returns:
            Optional[pd.Series]: 一个Pandas Series，代表了最新的那一行完整数据。
                                 如果历史数据为空（例如回测刚开始），则返回None。
        """
        if self._historical_data.empty:
            return None

        # .iloc[-1] 是从DataFrame中获取最后一行的最高效方式。
        return self._historical_data.iloc[-1]

    def get_history(self, bars: int) -> pd.DataFrame:
        """
        获取最近N条历史K线数据（及其所有指标）。

        当因子需要多于一个时间点的数据（例如计算交叉）时使用。

        Args:
            bars (int): 需要获取的历史K线的条数。

        Returns:
            pd.DataFrame: 一个Pandas DataFrame，包含了最近的N条历史数据。
                          如果可用的历史数据少于N条，则返回所有可用的数据。
                          如果历史数据为空，则返回一个空的DataFrame。
        """
        if self._historical_data.empty:
            return pd.DataFrame()

        # .tail(n) 是从DataFrame中获取最后N行的标准且高效的方式。
        return self._historical_data.tail(bars)


# --- 示例用法 ---
if __name__ == '__main__':
    # 1. 模拟一个已经准备好的 master_df
    data = {
        '15m_close': [100, 101, 102, 103, 104],
        '1h_close': [98.5, 98.5, 98.5, 102.5, 102.5],  # 经过ffill处理
        'ma_10': [99.8, 100.2, 100.8, 101.5, 102.3]
    }
    index = pd.to_datetime([
        '2025-09-12 10:00:00',
        '2025-09-12 10:15:00',
        '2025-09-12 10:30:00',
        '2025-09-12 10:45:00',
        '2025-09-12 11:00:00'
    ])
    master_df = pd.DataFrame(data, index=index)
    print("--- 完整的 Master DataFrame ---")
    print(master_df)
    print("\n" + "=" * 40 + "\n")

    # 2. 模拟回测引擎进行到 10:45:00 这一刻
    current_moment = pd.Timestamp('2025-09-12 10:45:00')
    print(f"--- 模拟引擎时间点: {current_moment} ---")

    # 3. 引擎为策略创建 Context 对象
    context = Context(master_df, current_moment)

    # 4. 策略通过 Context 获取数据
    print("\n--- 策略调用 context.get_latest_bar() ---")
    latest_bar = context.get_latest_bar()
    print("策略能看到的最新数据 (实际上是10:30的数据):")
    print(latest_bar)
    assert latest_bar.name == pd.Timestamp('2025-09-12 10:30:00')
    assert latest_bar['15m_close'] == 102

    print("\n--- 策略调用 context.get_history(2) ---")
    history = context.get_history(2)
    print("策略能看到的最近2条历史数据:")
    print(history)
    assert len(history) == 2
    assert history.index[-1] == pd.Timestamp('2025-09-12 10:30:00')
